///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/lanczos.h
///
/// \brief   header file for the lanczos iterative solver
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef FORMIC_LANCZOS_HEADER
#define FORMIC_LANCZOS_HEADER

#include <string>
#include <fstream>
#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <algorithm>

#include <boost/format.hpp>

#include <formic/numeric/numeric.h>
#include <formic/random/random.h>
#include <formic/lapack/interface.h>
#include <formic/mpi/interface.h>
#include <formic/timing/timing.h>
#include <formic/output/output.h>
#include <formic/matrix_actor.h>
#include <formic/exception.h>

namespace formic {

  template<class SCALAR>
  void lanczos_save_vec(const int n, const int i, const std::string & dir, const SCALAR * const vec) {
    //const std::string fname = ( boost::format("%slanczos_vec_%i.bin") % dir % i ).str();
    const std::string fname = ( boost::format("%s%i.bin") % dir % i ).str();
    std::ofstream ofs(fname.c_str(), std::ios::out | std::ios::binary);
    if (!ofs.is_open())
      throw formic::Exception( (boost::format("failed to open file \"%s\" in formic::lanczos_save_vec") % fname).str() );
    ofs.write((const char *)vec, size_t(n) * sizeof(SCALAR));
    ofs.close();
  }

  template<class SCALAR>
  void lanczos_load_vec(const int n, const int i, const std::string & dir, SCALAR * const vec) {
    //const std::string fname = ( boost::format("%slanczos_vec_%i.bin") % dir % i ).str();
    const std::string fname = ( boost::format("%s%i.bin") % dir % i ).str();
    std::ifstream ifs(fname.c_str(), std::ios::in | std::ios::binary);
    if (!ifs.is_open())
      throw formic::Exception( (boost::format("failed to open file \"%s\" in formic::lanczos_load_vec") % fname).str() );
    ifs.read((char *)vec, size_t(n) * sizeof(SCALAR));
    ifs.close();
  }

  //-----------------------------------------------------------------------------------
  // formic::lanczos -- Solves for the extremal eigenvalue and eigenvector of a matrix A.
  //-----------------------------------------------------------------------------------

  template<class SCALAR>
  inline bool lanczos(const bool get_evec,             // whether or not to compute the eigenvector
                      const int n,                     // dimension of the matrix
                      const int m,                     // maximum number of iterations
                      double & lambda,                 // size 1.                               on exit, the extremal eigenvalue.
                      SCALAR * const x,                // size n.                               on exit, the corresponding eigenvector if requested.
                      SCALAR * v0,                     // size n.    vector used as workspace.  on exit, contents are overwritten.
                      SCALAR * v1,                     // size n.    vector used as workspace.  on exit, contents are overwritten.
                      SCALAR * const t,                // size m*m.  vector used as workspace.  on exit, contents are overwritten.
                      SCALAR * const alpha,            // size m.    vector used as workspace.  on exit, contents are overwritten.
                      SCALAR * const beta,             // size m.    vector used as workspace.  on exit, contents are overwritten.
                      double * const w,                // size m.    vector used as workspace.  on exit, contents are overwritten.
                      SCALAR * const work,             // size 5*m.  vector used as workspace.  on exit, contents are overwritten.
                      double * const rwork,            // size 3*m.  vector used as workspace.  on exit, contents are overwritten.
                      formic::MatrixActor<SCALAR> * A, // object that applies linear transformations
                      const double thresh,             // energy convergence threshold
                      const bool restart = false,      // whether to restart from a saved vector
                      const double shift = 0.0,        // shift to apply to the matrix diagonal
                      const std::string & dir = "./",  // directory to use as scratch space, must end in a '/'
                      const int tid = 0)               // thread id
  {

    // start a timer
    formic::Stopwatch timer;
    timer.start();

    //// function object to save vectors to disk
    //struct SaveVec {
    //  void operator()(const int n, const int i, const std::string & dir, const SCALAR * const vec) {
    //    //const std::string fname = ( boost::format("%slanczos_vec_%i.bin") % dir % i ).str();
    //    const std::string fname = ( boost::format("%s%i.bin") % dir % i ).str();
    //    std::ofstream ofs(fname.c_str(), std::ios::out | std::ios::binary);
    //    if (!ofs.is_open())
    //      throw formic::Exception( (boost::format("failed to open file \"%s\" for writing in formic::lanczos") % fname).str() );
    //    ofs.write((const char *)vec, size_t(n) * sizeof(SCALAR));
    //    ofs.close();
    //  }
    //};
    //SaveVec save_vec;

    //// function object to load vectors from disk
    //struct LoadVec {
    //  void operator()(const int n, const int i, const std::string & dir, SCALAR * const vec) {
    //    //const std::string fname = ( boost::format("%slanczos_vec_%i.bin") % dir % i ).str();
    //    const std::string fname = ( boost::format("%s%i.bin") % dir % i ).str();
    //    std::ifstream ifs(fname.c_str(), std::ios::in | std::ios::binary);
    //    if (!ifs.is_open())
    //      throw formic::Exception( (boost::format("failed to open file \"%s\" for reading in formic::lanczos") % fname).str() );
    //    ifs.read((char *)vec, size_t(n) * sizeof(SCALAR));
    //    ifs.close();
    //  }
    //};
    //LoadVec load_vec;

    const std::string save_prefix = dir + "lanczos_vec_";
    const std::string solution_prefix = dir + "lanczos_solution_";

    //// get MPI info
    //const MPI::Comm & comm = MPI::COMM_WORLD;
    //const int nproc = comm.Get_size();
    //const int myrank = comm.Get_rank();

    const int nproc = 1;
    const int myrank = 0;

    // initialize the eigenvalue and the difference between the current and previous eigenvalues
    lambda = 0.0;
    double diff = 1.0e100;

    // either restart from a saved vector or choose a random, normalized vector for the first Krylov vector 
    if (tid == 0 && myrank == 0) {
      if (restart)
        formic::lanczos_load_vec(n, 0, solution_prefix, v1);
      else
        for (int i = 0; i < n; i++)
          v1[i] = formic::random_number<SCALAR>() - 0.5 * ( formic::unity(SCALAR()) + formic::imaginary_unity(SCALAR()) );
      formic::xscal(n, formic::unity(SCALAR()) / std::sqrt(formic::xdotc(n, v1, 1, v1, 1)), v1, 1);
      if (get_evec)
        formic::lanczos_save_vec(n, 0, save_prefix, v1);
    }
    //if (tid == 0) comm.Bcast(v1, n * sizeof(SCALAR), MPI::CHAR, 0);

    // print header for iteration printouts
    if (tid == 0 && myrank == 0) {
      //formic::of << boost::format("starting lanczos iterations at time %20.6f")
      //              % timer.elapsed_seconds()
      //           << std::endl;
      //formic::of << std::endl;
      formic::of << boost::format("        %6s     %20s     %20s     %20s")
                    % "iter"
                    % "smallest eigenvalue"
                    % "largest eigenvalue"
                    % "lanczos elapsed time"
                 << std::endl;
    }

    // iteratively enlarge the krylov subspace
    int iter = 0;
    for ( ; iter < m; iter++) {

      // check if we are finished 
      bool converged = ( std::abs(diff) < thresh );
      formic::mpi::bcast(converged);
      if (converged) break;

      // compute x = A * v1
      A->operate(n, v1, x, myrank, tid);

      // add shift
      if (tid == 0) formic::xaxpy(n, formic::unity(SCALAR()) * shift, v1, 1, x, 1);

      // orthogonalize against v0
      if (tid == 0 && iter > 0) formic::xaxpy(n, -beta[iter-1], v0, 1, x, 1);

      // compute alpha
      if (tid == 0) alpha[iter] = formic::xdotc(n, v1, 1, x, 1);

      // orthogonalize against v1
      if (tid == 0) formic::xaxpy(n, -alpha[iter], v1, 1, x, 1);

      // compute beta
      if (tid == 0) beta[iter] = std::sqrt(formic::xdotc(n, x, 1, x, 1));

      // normalize new krylov vector
      if (tid == 0) formic::xscal(n, formic::unity(SCALAR()) / beta[iter], x, 1);

      // save new krylov vector to disk
      if (get_evec && tid == 0 && myrank == 0) formic::lanczos_save_vec(n, iter+1, save_prefix, x);

      // prepare krylov vectors for next iteration
      std::swap(v0, v1);
      if (tid == 0) formic::xcopy(n, x, 1, v1, 1);

      // root node/thread solves eigenproblem in krylov basis
      if (tid == 0 && myrank == 0) {

        // construct A in the krylov basis
        const int dim = iter+1;
        formic::xscal(dim*dim, formic::zero(SCALAR()), t, 1);
        for (int i = 0; i < dim; i++) {
          t[dim*i+i] = alpha[i];
          if (i > 0)
            t[dim*(i-1)+i] = t[dim*i+(i-1)] = beta[i-1];
        }

        // diagonalize
        int info;
        int lwork = 5*dim;
        formic::xsyev('V', 'U', dim, t, dim, w, work, lwork, rwork, info);

        // compute eigenvalue difference and update the eigenvalue (accounting for the shift)
        diff = w[0] - shift - lambda;
        lambda = w[0] - shift;

        // print eigenvalue for this iteration
        formic::of << boost::format("        %6i     %20.12f     %20.12f     %20.6f")
                      % iter
                      % lambda
                      % (w[dim-1]-shift)
                      % timer.elapsed_seconds()
                   << std::endl;

        //for (int i = 0; i < dim; i++)
        //  formic::of << boost::format(" %16.4e") % w[i];
        //formic::of << std::endl;

      }

    }

    // check if the method converged
    bool converged = ( std::abs(diff) < thresh );
    formic::mpi::bcast(converged);
    if (converged && myrank == 0 && tid == 0)
      formic::of << boost::format("lanczos converged in %10i iterations") % iter << std::endl << std::endl;
    else if (myrank == 0 && tid == 0)
      formic::of << boost::format("lanczos did not converge") << std::endl << std::endl;

    if (get_evec && tid == 0 && myrank == 0) {

      // load vectors from disk and construct solution vector
      formic::xscal(n, formic::zero(SCALAR()), x, 1);
      for (int i = 0; i < iter; i++) {
        formic::lanczos_load_vec(n, i, save_prefix, v0);
        formic::xaxpy(n, t[i], v0, 1, x, 1);
      }

      // save solution vector to disk
      formic::lanczos_save_vec(n, 0, solution_prefix, x);

    }

    //if (tid == 0) comm.Bcast(x, n * sizeof(SCALAR), MPI::CHAR, 0);

    //if (tid == 0) comm.Bcast(&lambda, sizeof(double), MPI::CHAR, 0);

    return converged;

  }

} // end namespace formic

#endif
